from flask import Flask, render_template, request, redirect, url_for, flash, jsonify
from flask_sqlalchemy import SQLAlchemy
from werkzeug.security import generate_password_hash, check_password_hash
import os
import secrets
from datetime import datetime, timedelta

app = Flask(__name__)
app.config['SECRET_KEY'] = os.environ.get('SECRET_KEY', 'dev_key_12345')
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///messages.db'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False

db = SQLAlchemy(app)

# 消息模型
class Message(db.Model):
    id = db.Column(db.String(16), primary_key=True)
    content = db.Column(db.Text, nullable=False)
    password_hash = db.Column(db.String(128), nullable=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    expires_at = db.Column(db.DateTime, default=lambda: datetime.utcnow() + timedelta(days=7))
    locker_id = db.Column(db.Integer, db.ForeignKey('locker.id'), nullable=True)

    def __repr__(self):
        return f'<Message {self.id}>'

# 消息柜模型
class Locker(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(50), nullable=False)
    capacity = db.Column(db.Integer, nullable=False, default=10)
    messages = db.relationship('Message', backref='locker', lazy=True)
    
    def get_used_space(self):
        return len(self.messages)
    
    def is_full(self):
        return self.get_used_space() >= self.capacity
    
    def __repr__(self):
        return f'<Locker {self.id}: {self.name}>'

# 创建数据库表并初始化消息柜
with app.app_context():
    db.create_all()
    
    # 初始化消息柜
    if not Locker.query.first():
        print("初始化消息柜...")
        lockers = [
            Locker(name=f"消息柜 #{i}", capacity=10) for i in range(1, 11)
        ]
        db.session.add_all(lockers)
        db.session.commit()
        print(f"成功创建 {len(lockers)} 个消息柜")

# 路由
@app.route('/')
def index():
    return render_template('index.html')

@app.route('/lockers')
def lockers():
    all_lockers = Locker.query.all()
    return render_template('lockers.html', lockers=all_lockers)

@app.route('/locker/<int:locker_id>')
def locker_detail(locker_id):
    locker = Locker.query.get_or_404(locker_id)
    return render_template('locker_detail.html', locker=locker)

@app.route('/create', methods=['POST'])
def create_message():
    content = request.form.get('content')
    password = request.form.get('password')
    
    # 验证输入
    if not content or not password:
        return jsonify({'success': False, 'error': '消息内容和密码都是必填项'}), 400
    
    if not password.isdigit() or len(password) != 6:
        return jsonify({'success': False, 'error': '密码必须是6位数字'}), 400
    
    # 生成唯一ID
    message_id = secrets.token_urlsafe(12)
    
    # 创建新消息
    new_message = Message(
        id=message_id,
        content=content,
        password_hash=generate_password_hash(password)
    )
    
    # 保存到数据库
    db.session.add(new_message)
    db.session.commit()
    
    # 返回成功信息
    return jsonify({
        'success': True,
        'message_id': message_id,
        'link': url_for('index', _external=True) + '#' + message_id
    })

@app.route('/create_in_locker', methods=['POST'])
def create_in_locker():
    content = request.form.get('content')
    password = request.form.get('password')
    locker_id = request.form.get('locker_id')
    
    # 验证输入
    if not content or not password or not locker_id:
        return jsonify({'success': False, 'error': '内容、密码和柜子ID不能为空'}), 400
    
    if not password.isdigit() or len(password) != 6:
        return jsonify({'success': False, 'error': '密码必须是6位数字'}), 400
    
    locker = Locker.query.get(locker_id)
    if not locker:
        return jsonify({'success': False, 'error': '找不到指定的消息柜'}), 404
    
    if locker.is_full():
        return jsonify({'success': False, 'error': '该消息柜已满，请选择其他柜子'}), 400
    
    # 生成唯一ID
    message_id = secrets.token_urlsafe(12)
    
    # 创建新消息
    new_message = Message(
        id=message_id,
        content=content,
        password_hash=generate_password_hash(password),
        locker_id=locker_id
    )
    
    # 保存到数据库
    db.session.add(new_message)
    db.session.commit()
    
    # 返回成功信息
    return jsonify({
        'success': True,
        'message_id': message_id,
        'locker_id': locker_id,
        'locker_name': locker.name
    })

@app.route('/retrieve', methods=['POST'])
def retrieve_message():
    message_id = request.form.get('message_id')
    password = request.form.get('password')
    
    # 验证输入
    if not message_id or not password:
        return jsonify({'success': False, 'error': '消息ID和密码都是必填项'}), 400
    
    # 查找消息
    message = Message.query.get(message_id)
    if not message:
        return jsonify({'success': False, 'error': '找不到指定的消息'}), 404
    
    # 检查密码
    if not check_password_hash(message.password_hash, password):
        return jsonify({'success': False, 'error': '密码不正确'}), 403
    
    # 获取消息内容
    content = message.content
    
    # 删除消息
    db.session.delete(message)
    db.session.commit()
    
    # 返回消息内容
    return jsonify({'success': True, 'content': content})

@app.route('/retrieve_from_locker', methods=['POST'])
def retrieve_from_locker():
    locker_id = request.form.get('locker_id')
    password = request.form.get('password')
    
    # 验证输入
    if not locker_id or not password:
        return jsonify({'success': False, 'error': '柜子ID和密码不能为空'}), 400
    
    locker = Locker.query.get(locker_id)
    if not locker:
        return jsonify({'success': False, 'error': '找不到指定的消息柜'}), 404
    
    # 查找该柜子中所有消息
    messages = Message.query.filter_by(locker_id=locker_id).all()
    
    if not messages:
        return jsonify({'success': False, 'error': '该消息柜中没有消息'}), 404
    
    # 尝试用密码解锁每一条消息
    for message in messages:
        if check_password_hash(message.password_hash, password):
            content = message.content
            
            # 删除消息
            db.session.delete(message)
            db.session.commit()
            
            return jsonify({
                'success': True,
                'content': content,
                'locker_name': locker.name
            })
    
    return jsonify({'success': False, 'error': '密码错误或没有匹配的消息'}), 403

@app.route('/cleanup')
def cleanup_expired():
    """管理员路由，清理过期消息"""
    now = datetime.utcnow()
    expired = Message.query.filter(Message.expires_at < now).all()
    count = len(expired)
    
    for message in expired:
        db.session.delete(message)
    
    db.session.commit()
    return f"已清理 {count} 条过期消息"

if __name__ == '__main__':
    app.run(debug=True, port=5001)
